#!/bin/bash -x

test_pod_dir="$HOME/workspace/pod-workload"

base_image_name="busybox:latest"
base_image_uri=""
unencrypted_image_name="busybox-unencrypted:latest"
unencrypted_image_uri=""
encrypted_image_name="busybox-encrypted:latest"
encrypted_image_uri=""

registry_ip_addr=""
registry_port=31313
registry=""

pod_yaml="${test_pod_dir}/encrypted-image-pod.yaml"
pod_name="busybox"
pod_runtimeclass="kata-qemu-csv"
pod_container_image_uri=""

coco_keyprovider_image_uri="docker.io/robertsonhan/coco-keyprovider:v0.10.0"

usage () {
  set +x
  echo -e "\033[31m
Usage: $0 <-h|<-i <registry-ip-addr>> [-p <registry-port>] [-r <kata-coco-runtimeclass>] [-C <coco-keyprovider-image-uri>]>

       -h,--help            Display this help message
       -i,--ip              The ip address of the local registry
       -p,--port            The port number of the local registry
       -r,--runtimeclass    Specify the runtimeclass
                            Valid value are shown as following:
                                kata-qemu-coco-dev
                                kata-qemu-csv
                                kata-qemu-csv2
                                kata-qemu-csv3
                            Default is kata-qemu-csv
       -C,--cocokeyprovider The uri of CoCo KeyProvider payload image
                            Default is docker.io/robertsonhan/coco-keyprovider:v0.10.0
\033[0m"
  set -x
}

params=$(getopt -o hi:p:r:C: --long help,ip:,port:,runtimeclass:,cocokeyprovider: -- "$@")
if [ $? -ne 0 ]; then
  echo "[Error] cannot parse parameters"
  exit 1
fi

eval set -- "$params"

while true; do
  case "$1" in
  -h|--help)
    usage
    exit 0
    ;;
  -i|--ip)
    registry_ip_addr="$2"
    shift 2
    ;;
  -p|--port)
    registry_port="$2"
    shift 2
    ;;
  -r|--runtimeclass)
    pod_runtimeclass="$2"
    shift 2
    ;;
  -C|--cocokeyprovider)
    coco_keyprovider_image_uri="$2"
    shift 2
    ;;
  --)
    shift
    break
    ;;
  *)
    usage
    exit 1
    ;;
  esac
done

if [ "X${registry_ip_addr}" = "X" ]; then
  usage
  exit 1
fi

registry="${registry_ip_addr}:${registry_port}"
base_image_uri="${registry}/${base_image_name}"
unencrypted_image_uri="${registry}/${unencrypted_image_name}"
encrypted_image_uri="${registry}/${encrypted_image_name}"
pod_container_image_uri="${encrypted_image_uri}"



CONFIDENTIAL_SECRET="test image's encryption secret"
skopeo_bin=$(which skopeo)

kbs_namespace="trustee-operator-system"
kbs_deployment="trustee-deployment"
kbs_ip_addr="$(kubectl get pods -n ${kbs_namespace} -o wide | \
  grep ${kbs_deployment} | \
  sed "s/^.* \([0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\) .*$/\1/g")"
kbs_resource_path="/opt/confidential-containers/kbs/repository"
kbs_key_path="default/image_key/${pod_name}"
kbs_key_id="kbs:///${kbs_key_path}"

image_enc_key_file="${test_pod_dir}/image_enc_key"
image_enc_key_b64=""

generate_image_enc_key () {
  echo "[Generate image encryption key]"
  if [ ! -e ${image_enc_key_file} ]; then
    head -c 32 /dev/urandom | openssl enc > ${image_enc_key_file}
  fi

  echo "[  Convert image encryption key to base64-encoded format]"
  image_enc_key_b64="$(base64 < ${image_enc_key_file})"
}

generate_image_enc_key

upload_image_encryption_key_to_kbs () {
  echo "[Upload image encryption key to kbs]"
  kubectl exec deploy/${kbs_deployment} -n ${kbs_namespace} \
    -c kbs \
    -- mkdir -p "${kbs_resource_path}/$(dirname "${kbs_key_path}")"
  cat "${image_enc_key_file}" | \
    kubectl exec -i deploy/${kbs_deployment} -n ${kbs_namespace} \
    -c kbs \
    -- tee "${kbs_resource_path}/${kbs_key_path}" > /dev/null
}

upload_image_encryption_key_to_kbs

build_encrypted_image () {
  echo "[Build unencrypted & encrypted image based on base image]"

  echo "[  Build unencrypted image based on base image]"
  local docker_build_arg_https_proxy=""
  local docker_build_arg_http_proxy=""

  if [ "X${https_proxy}" != "X" ]; then
    docker_build_arg_https_proxy="--build-arg HTTPS_PROXY=${https_proxy}"
  elif [ "X${HTTPS_PROXY}" != "X" ]; then
    docker_build_arg_https_proxy="--build-arg HTTPS_PROXY=${HTTPS_PROXY}"
  fi
  if [ "X${http_proxy}" != "X" ]; then
    docker_build_arg_http_proxy="--build-arg HTTP_PROXY=${http_proxy}"
  elif [ "X${HTTP_PROXY}" != "X" ]; then
    docker_build_arg_http_proxy="--build-arg HTTP_PROXY=${HTTP_PROXY}"
  fi

  docker build \
    ${docker_build_arg_https_proxy} \
    ${docker_build_arg_http_proxy} \
    -t ${unencrypted_image_uri} - <<EOF
FROM ${base_image_uri}
RUN echo "${CONFIDENTIAL_SECRET}" > /secret
EOF

  echo "[  Build encrypted image based on unencrypted image]"
  local skopeo_input_dir="${test_pod_dir}/image_encryption/oci/input"
  local skopeo_output_dir="${test_pod_dir}/image_encryption/oci/output"

  if [ ! -d ${skopeo_input_dir} ]; then
    mkdir -p ${skopeo_input_dir} || true
  fi
  rm -rf ${skopeo_input_dir}/* || true
  if [ ! -d ${skopeo_output_dir} ]; then
    mkdir -p ${skopeo_output_dir} || true
  fi
  rm -rf ${skopeo_output_dir}/* || true

  echo "[    Use skopeo to provide image info of the unencrypted image]"
  ${skopeo_bin} --insecure-policy \
    copy docker-daemon:${unencrypted_image_uri} \
    dir:${skopeo_input_dir}

  echo "[    Use CoCo KeyProvider to encrypt the unencrypted image]"
  docker run \
    -v "${test_pod_dir}/image_encryption/oci:/oci" \
    ${coco_keyprovider_image_uri} \
    /encrypt.sh \
    -k "${image_enc_key_b64}" \
    -i "${kbs_key_id}" \
    -s dir:/oci/input -d dir:/oci/output

  command jq --version > /dev/null
  if [ $? -ne 0 ]; then
    sudo yum install -y jq
  fi
  ${skopeo_bin} inspect dir:${skopeo_output_dir} | jq \
    '.LayersData[0].Annotations["org.opencontainers.image.enc.keys.provider.attestation-agent"] | @base64d | fromjson'

  echo "[    Upload encrypted image to the registry]"
  ${skopeo_bin} --insecure-policy \
    copy dir:${skopeo_output_dir} "docker://${encrypted_image_uri}"
}

build_encrypted_image

generate_pod_yaml () {
  echo "[Generate Pod yaml]"

  if [ ! -d $(dirname ${pod_yaml}) ]; then
    mkdir -p $(dirname ${pod_yaml})
  fi

  cat <<EOF | tee ${pod_yaml}
apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: ${pod_name}
  name: encrypted-${pod_name}
spec:
  replicas: 1
  selector:
    matchLabels:
      app: ${pod_name}
  template:
    metadata:
      labels:
        app: ${pod_name}
      annotations:
        io.katacontainers.config.hypervisor.kernel_params: 'agent.guest_components_rest_api=resource agent.aa_kbc_params=cc_kbc::http://${kbs_ip_addr}:8080'
        io.containerd.cri.runtime-handler: ${pod_runtimeclass}
    spec:
      runtimeClassName: ${pod_runtimeclass}
      containers:
      - image: ${pod_container_image_uri}
        name: ${pod_name}-c1
        imagePullPolicy: Always
EOF
}

generate_pod_yaml

run_pod () {
  echo "[Run Pod]"
  kubectl delete -f ${pod_yaml} || true
  kubectl apply -f ${pod_yaml}
}

run_pod
